#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <nav_msgs/Path.h>
#include <ompl/tools/benchmark/Benchmark.h>
#include <ompl/control/planners/rrt/RRT.h>
#include <ompl/control/planners/kpiece/KPIECE1.h>
#include <ssrr_ompl_planner/SsrrPlanning.h>
#include <boost/math/constants/constants.hpp>
#include <bullet_rviz_test/stability_checker.h>
#include <tf/tf.h>
#include <tf/transform_listener.h>
#include <std_msgs/Bool.h>

using namespace ssrr;

ros::Publisher path_pub_;
std::string global_frame_, base_frame_;
bullet_rviz_test::StabilityChecker* checker_;
geometry_msgs::PoseStamped robot_pose_;
bool make_plan_;

void retrieveValues(const ompl::base::State *state, double &x, double &y, double &th) {
    const ompl::base::SE2StateSpace::StateType* se2State = state->as<ompl::base::SE2StateSpace::StateType>();
    const ompl::base::RealVectorStateSpace::StateType* xy = se2State->as<ompl::base::RealVectorStateSpace::StateType>(0);
    const ompl::base::SO2StateSpace::StateType* so2 = se2State->as<ompl::base::SO2StateSpace::StateType>(1);
    x = xy->values[0];
    y = xy->values[1];
    th = so2->value;
}

class SSRRStateSampler : public ompl::base::ValidStateSampler
{
public:
    SSRRStateSampler(const ompl::base::SpaceInformation *si) : ValidStateSampler(si)
    {
        name_ = "dummy ssrr sampler";
    }
    // Generate a sample in the valid part of the R^3 state space
    // Valid states satisfy the following constraints:
    // -1<= x,y,z <=1
    // if .25 <= z <= .5, then |x|>.8 and |y|>.8
    virtual bool sample(ompl::base::State *state)
    {
        double* val = static_cast<ompl::base::RealVectorStateSpace::StateType*>(state)->values;
        double z = rng_.uniformReal(-1,1);

        if (z>.25 && z<.5)
        {
            double x = rng_.uniformReal(0,1.8), y = rng_.uniformReal(0,.2);
            switch(rng_.uniformInt(0,3))
            {
            case 0: val[0]=x-1;  val[1]=y-1;
                break;
            case 1: val[0]=x-.8; val[1]=y+.8;
                break;
            case 2: val[0]=y-1;  val[1]=x-1;
                break;
            case 3: val[0]=y+.8; val[1]=x-.8;
                break;
            }
        }
        else
        {
            val[0] = rng_.uniformReal(-1,1);
            val[1] = rng_.uniformReal(-1,1);
        }
        val[2] = z;
        std::cout<< "return sample " << val[0] << " " << val[1] << " " << val[2] << std::endl;
        //assert(si_->isValid(state));
        return true;
    }
    // We don't need this in the example below.
    virtual bool sampleNear(ompl::base::State*, const ompl::base::State*, const double)
    {
        throw ompl::Exception("MyValidStateSampler::sampleNear", "not implemented");
        return false;
    }
protected:
    ompl::RNG rng_;
};

// return an instance of my sampler
ompl::base::ValidStateSamplerPtr allocSSRRStateSampler(const ompl::base::SpaceInformation *si)
{
    return ompl::base::ValidStateSamplerPtr(new SSRRStateSampler(si));
}

bool isValid(const ompl::base::State* state) {
    const double* val = static_cast<const ompl::base::RealVectorStateSpace::StateType*>(state)->values;
    double x = val[0];
    double y = val[1];
    double th = val[2];

    bool res = checker_->check(x, y, th);

    if (res)
        return true;
    else
    {
        ROS_INFO("State not valid!");
        return false;
    }

}

void getRobotPosition(geometry_msgs::PoseStamped& pose_stamped, tf::TransformListener& tf_){


    // pose
    tf::Stamped<tf::Pose> robot_pose;
    robot_pose.setIdentity();
    robot_pose.frame_id_ = base_frame_;
    robot_pose.stamp_ = ros::Time::now();

    // ROBOT POSE PUBLISHER

    tf_.waitForTransform(global_frame_, base_frame_, ros::Time(), ros::Duration(1.0));
    tf::StampedTransform transform;
    try
    {
        tf_.lookupTransform(global_frame_, base_frame_, ros::Time(0), transform);

        // construct a pose message
        pose_stamped.header.frame_id = global_frame_;
        pose_stamped.header.stamp = ros::Time::now();

        pose_stamped.pose.orientation.x = transform.getRotation().getX();
        pose_stamped.pose.orientation.y = transform.getRotation().getY();
        pose_stamped.pose.orientation.z = transform.getRotation().getZ();
        pose_stamped.pose.orientation.w = transform.getRotation().getW();

        pose_stamped.pose.position.x = transform.getOrigin().getX();
        pose_stamped.pose.position.y = transform.getOrigin().getY();
        pose_stamped.pose.position.z = transform.getOrigin().getZ();

    }
    catch (tf::TransformException &ex)
    {
        ROS_ERROR("SSRR_PLANNING: %s",ex.what());
    }
}

void startup(SsrrPlanning& setup){

    // plan for kinematic car in SE(2)
    ompl::base::StateSpacePtr SE2(setup.getStateSpace());

    // set the bounds for the R^2 part of SE(2)
    ompl::base::RealVectorBounds bounds(2);
    bounds.setLow(-4.5); // derive from map
    bounds.setHigh(2.5);
    SE2->as<ompl::base::SE2StateSpace>()->setBounds(bounds);

    // define start state
    ompl::base::ScopedState<ompl::base::SE2StateSpace> start(SE2);

    start->setX(robot_pose_.pose.position.x);
    start->setY(robot_pose_.pose.position.y);
    //tf::Quaternion q = tf::createQuaternionFromYaw(locations_[index].orientation);
    double yaw = tf::getYaw(robot_pose_.pose.orientation);
    start->setYaw(yaw);

    // define goal state
    ompl::base::ScopedState<ompl::base::SE2StateSpace> goal(SE2);
    goal->setX(-3);
    goal->setY(-3);
    goal->setYaw(M_PI);

    // set the start & goal states
    setup.setStartAndGoalStates(start, goal, .1);

    //state sampler
    //setup.getSpaceInformation()->setValidStateSamplerAllocator(allocSSRRStateSampler);

}

void publishPlan(ompl::control::PathControl& path){
    uint number_of_states = path.getStateCount();
    std::vector<ompl::base::State*> states = path.getStates();
    ROS_INFO("Path as geometric: number_of_states: %d", number_of_states);

    nav_msgs::Path ompl_plan;
    ompl_plan.header.frame_id = global_frame_;
    ompl_plan.header.stamp = ros::Time::now();

    for(uint index=0; index < number_of_states; index++){

        geometry_msgs::PoseStamped pose;
        pose.header.frame_id = global_frame_;
        pose.header.stamp = ros::Time::now();

        double x, y, th;
        retrieveValues(states[index], x, y, th);
        pose.pose.position.x = x;
        pose.pose.position.y = y;
        pose.pose.position.z = 0;
        ROS_INFO("state[%d]: x[%.1f] y[%.1f] th[%.1f]", index, x, y, th);

        tf::Quaternion q = tf::createQuaternionFromYaw(th);
        tf::quaternionTFToMsg(q, pose.pose.orientation);

        ompl_plan.poses.push_back(pose);
    }

    path_pub_.publish(ompl_plan);
}

void plan()
{

    SsrrPlanning setup;

    //init start, goal state and state sampler
    startup(setup);

    //Plan
    setup.setPlanner(ompl::base::PlannerPtr(new ompl::control::RRT(setup.getSpaceInformation())));
    std::vector<double> cs(2);
    cs[0] = cs[1] = 0.1;
    setup.setup();
    setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);

    // try to solve the problem
    if (setup.solve(20))
    {
        // print the (approximate) solution path: print states along the path
        // and controls required to get from one state to the next
        ompl::control::PathControl& path(setup.getSolutionPath());
        //path.interpolate(); // uncomment if you want to plot the path
        //path.print(std::cout);

        publishPlan(path);

        if (!setup.haveExactSolutionPath())
        {
            std::cout << "Solution is approximate. Distance to actual goal is " <<
                         setup.getProblemDefinition()->getSolutionDifference() << std::endl;
        }
    }
}

void makePlan(std_msgs::Bool msg){

    make_plan_ = msg.data;
    ROS_INFO("plan cb received : %d", make_plan_);
}

int main(int argc, char *argv[])
{
    ROS_INFO("SSRR_Planing_NODE");

    ros::init(argc, argv, "ssrr_planning_node");
    ros::NodeHandle nh;
    ros::NodeHandle private_nh("~");

    tf::TransformListener tf_;
    base_frame_ = "/base_link";
    global_frame_ = "/map";

    make_plan_ = false;

    path_pub_ = private_nh.advertise<nav_msgs::Path>("ompl_plan", 10);
    checker_ = new bullet_rviz_test::StabilityChecker(false, false);

    ros::Subscriber makePlanSub = private_nh.subscribe<std_msgs::Bool>("make_plan", 10, makePlan);

    ros::Rate freq = ros::Rate(1);

    while(nh.ok()){

        //if (checker_->has_map() && make_plan_)
        if (make_plan_)
        {
            ROS_INFO("making a plan");
            getRobotPosition(robot_pose_, tf_);
            plan();
            make_plan_ = false;
        }

        ros::spinOnce();
        freq.sleep();
    }
}
